// test1.rs
use tokio::time::{Duration, Instant};
use zeromq::{prelude::*, *};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

static TOTAL_MESSAGES: AtomicU64 = AtomicU64::new(0);

async fn start_server_01() -> Result<(), Box<dyn std::error::Error>> {
    // 创建 REP socket 用于接收消息
    let mut frontend = zeromq::RepSocket::new();
    frontend.bind("tcp://127.0.0.1:5555").await?;

    // 创建 PUB socket 用于发布消息
    let mut backend = zeromq::PubSocket::new();
    backend.bind("tcp://127.0.0.1:5556").await?;

    println!("消息代理已启动: 接收端口 5555, 发布端口 5556");

    loop {
        // 接收消息
        let msg = frontend.recv().await?;

        // 转发消息到发布者
        backend.send(msg).await?;

        // 发送确认回复
        let reply = zeromq::ZmqMessage::from("OK");
        frontend.send(reply).await?;
    }
}

async fn start_producer_01(id: usize) -> Result<(), Box<dyn std::error::Error>> {
    let mut socket = zeromq::ReqSocket::new();
    socket.connect("tcp://127.0.0.1:5555").await?;

    const MESSAGE_COUNT: i32 = 100_000; // 每个生产者发送1000条消息

    for i in 0..MESSAGE_COUNT {
        let payload = format!("Producer-{}-Message-{}", id, i);
        let msg = zeromq::ZmqMessage::from(payload);

        // 发送消息
        socket.send(msg).await?;

        // 等待确认
        let _reply = socket.recv().await?;
    }

    println!("生产者 {} 完成发送 {} 条消息", id, MESSAGE_COUNT);
    Ok(())
}

async fn start_consumer_01(id: usize) -> Result<(), Box<dyn std::error::Error>> {
    let mut socket = zeromq::SubSocket::new();
    socket.connect("tcp://127.0.0.1:5556").await?;

    // 订阅所有消息
    socket.subscribe("").await?;

    println!("消费者 {} 已连接并开始接收消息", id);

    let mut count = 0u64;
    const MAX_MESSAGES: u64 = 1_000_000; // 每个消费者最多接收10000条消息

    while count < MAX_MESSAGES {
        match tokio::time::timeout(Duration::from_secs(5), socket.recv()).await {
            Ok(Ok(_msg)) => {
                count += 1;
                TOTAL_MESSAGES.fetch_add(1, Ordering::Relaxed);

                // 每1000条消息打印一次进度
                if count % 1000 == 0 {
                    println!("消费者 {} 已接收 {} 条消息", id, count);
                }
            }
            Ok(Err(e)) => {
                eprintln!("消费者 {} 接收错误: {}", id, e);
                break;
            }
            Err(_) => {
                // 超时，可能没有更多消息
                println!("消费者 {} 超时，停止接收", id);
                break;
            }
        }
    }

    println!("消费者 {} 完成接收 {} 条消息", id, count);
    Ok(())
}

pub async fn run_performance_test_01() -> Result<(), Box<dyn std::error::Error>> {
    // 启动服务端（消息代理）
    tokio::spawn(async {
        if let Err(e) = start_server_01().await {
            eprintln!("Server error: {}", e);
        }
    });

    // 等待服务器启动
    tokio::time::sleep(Duration::from_millis(100)).await;

    // 创建生产者和消费者
    let producer_count = 10;
    let consumer_count = 10;

    let mut handles = vec![];

    // 启动消费者
    for i in 0..consumer_count {
        let handle = tokio::spawn(async move {
            if let Err(e) = start_consumer_01(i).await {
                eprintln!("Consumer {} error: {}", i, e);
            }
        });
        handles.push(handle);
    }

    // 等待消费者连接
    tokio::time::sleep(Duration::from_millis(100)).await;

    // 启动生产者
    let start_time = Instant::now();
    for i in 0..producer_count {
        let handle = tokio::spawn(async move {
            if let Err(e) = start_producer_01(i).await {
                eprintln!("Producer {} error: {}", i, e);
            }
        });
        handles.push(handle);
    }

    // 等待所有任务完成
    for handle in handles {
        handle.await?;
    }

    let duration = start_time.elapsed();
    let total_messages = TOTAL_MESSAGES.load(Ordering::Relaxed);

    println!("测试完成:");
    println!("总消息数: {}", total_messages);
    println!("耗时: {:?}", duration);
    println!("吞吐量: {:.2} 消息/秒", total_messages as f64 / duration.as_secs_f64());
    
    Ok(())
}